{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "62306894",
   "metadata": {},
   "source": [
    "# 1. Environment Setup\n",
    "Seed the random generators, import core dependencies, and detect the training device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b78ed595",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on device: cpu\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from pyhealth.datasets import SleepEDFDataset\n",
    "from pyhealth.datasets.splitter import split_by_sample\n",
    "from pyhealth.datasets.utils import get_dataloader\n",
    "\n",
    "SEED = 42\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(SEED)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Running on device: {device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eba8ce8",
   "metadata": {},
   "source": [
    "# 2. Load Sleep-EDF Dataset\n",
    "Point to the Sleep-EDF dataset root and load the telemetry subset for sleep stage classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "83784e4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No config path provided, using default config\n",
      "Initializing sleepedf dataset from ../downloads/sleepedf (dev mode: False)\n",
      "Scanning table: recordings from F:\\coding_projects\\pyhealth\\downloads\\sleepedf\\sleepedf-cassette-pyhealth.csv\n",
      "Collecting global event dataframe...\n",
      "Collected dataframe with shape: (2, 9)\n",
      "Dataset: sleepedf\n",
      "Dev mode: False\n",
      "Number of patients: 1\n",
      "Number of events: 2\n"
     ]
    }
   ],
   "source": [
    "dataset = SleepEDFDataset(\n",
    "    root=\"../downloads/sleepedf\",  # Update this path\n",
    ")\n",
    "dataset.stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cc3085e",
   "metadata": {},
   "source": [
    "# 3. Prepare PyHealth Dataset\n",
    "Set the task for the dataset and convert raw samples into PyHealth format for self-supervised learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2e8c14b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting task SleepStaging for sleepedf base dataset...\n",
      "Generating samples with 1 worker(s)...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating samples for SleepStaging with 1 worker:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Used Annotations descriptions: ['Sleep stage 1', 'Sleep stage 2', 'Sleep stage 3', 'Sleep stage 4', 'Sleep stage R', 'Sleep stage W']\n",
      "Not setting metadata\n",
      "2650 matching events found\n",
      "No baseline correction applied\n",
      "0 projection items activated\n",
      "Using data from preloaded Raw for 2650 events and 3000 original time points ...\n",
      "0 bad epochs dropped\n",
      "Used Annotations descriptions: ['Sleep stage 1', 'Sleep stage 2', 'Sleep stage 3', 'Sleep stage 4', 'Sleep stage R', 'Sleep stage W']\n",
      "Not setting metadata\n",
      "2829 matching events found\n",
      "No baseline correction applied\n",
      "0 projection items activated\n",
      "Using data from preloaded Raw for 2829 events and 3000 original time points ...\n",
      "0 bad epochs dropped\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating samples for SleepStaging with 1 worker: 100%|██████████| 1/1 [00:14<00:00, 14.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Label label vocab: {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Processing samples: 100%|██████████| 5479/5479 [00:00<00:00, 11703.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 5479 samples for task SleepStaging\n",
      "Total task samples: 5479\n",
      "Input schema: {'signal': 'tensor'}\n",
      "Output schema: {'label': 'multiclass'}\n",
      "\n",
      "Sample keys: dict_keys(['patient_id', 'night', 'patient_age', 'patient_sex', 'signal', 'label'])\n",
      "Signal shape: torch.Size([7, 3000])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "sample_dataset = dataset.set_task()\n",
    "\n",
    "print(f\"Total task samples: {len(sample_dataset)}\")\n",
    "print(f\"Input schema: {sample_dataset.input_schema}\")\n",
    "print(f\"Output schema: {sample_dataset.output_schema}\")\n",
    "\n",
    "if len(sample_dataset) == 0:\n",
    "    raise RuntimeError(\"The task did not produce any samples. Verify the dataset root.\")\n",
    "\n",
    "# Inspect a sample\n",
    "sample = sample_dataset[0]\n",
    "print(f\"\\nSample keys: {sample.keys()}\")\n",
    "print(f\"Signal shape: {sample['signal'].shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4789868f",
   "metadata": {},
   "source": [
    "# 4. Split Dataset\n",
    "Divide the processed samples into training, validation, and test subsets before building dataloaders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a34eb15b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train/Val/Test sizes: 3835, 548, 1096\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = 32\n",
    "\n",
    "train_ds, val_ds, test_ds = split_by_sample(sample_dataset, [0.7, 0.1, 0.2], seed=SEED)\n",
    "print(f\"Train/Val/Test sizes: {len(train_ds)}, {len(val_ds)}, {len(test_ds)}\")\n",
    "\n",
    "train_loader = get_dataloader(train_ds, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_loader = get_dataloader(val_ds, batch_size=BATCH_SIZE) if len(val_ds) else None\n",
    "test_loader = get_dataloader(test_ds, batch_size=BATCH_SIZE) if len(test_ds) else None\n",
    "\n",
    "if len(train_loader) == 0:\n",
    "    raise RuntimeError(\"The training loader is empty. Increase the dataset size or adjust the split ratios.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c85a3195",
   "metadata": {},
   "source": [
    "# 5. Inspect Batch Structure\n",
    "Peek at the first training batch to understand feature shapes and data structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "99fa15e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch structure:\n",
      "  patient_id: list(len=32)\n",
      "  night: list(len=32)\n",
      "  patient_age: list(len=32)\n",
      "  patient_sex: list(len=32)\n",
      "  signal: Tensor(shape=(32, 7, 3000))\n",
      "  label: Tensor(shape=(32,))\n"
     ]
    }
   ],
   "source": [
    "first_batch = next(iter(train_loader))\n",
    "\n",
    "def describe(value):\n",
    "    if hasattr(value, \"shape\"):\n",
    "        return f\"{type(value).__name__}(shape={tuple(value.shape)})\"\n",
    "    if isinstance(value, (list, tuple)):\n",
    "        return f\"{type(value).__name__}(len={len(value)})\"\n",
    "    return type(value).__name__\n",
    "\n",
    "batch_summary = {key: describe(value) for key, value in first_batch.items()}\n",
    "print(\"Batch structure:\")\n",
    "for key, desc in batch_summary.items():\n",
    "    print(f\"  {key}: {desc}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faf8366e",
   "metadata": {},
   "source": [
    "# 6. Instantiate ContraWR Model\n",
    "Create the PyHealth ContraWR model for self-supervised learning on sleep signals and review its parameter footprint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6ba3aff5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "f:\\coding_projects\\pyhealth\\.venv\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Input data dimensions ===\n",
      "n_channels: 7\n",
      "length: 3000\n",
      "=== Spectrogram parameters ===\n",
      "n_channels: 7\n",
      "freq_dim: 65\n",
      "time_steps: 90\n",
      "=== Convolution Parameters ===\n",
      "in_channels: 7, out_channels: 8, freq_dim: 16, time_steps: 22\n",
      "in_channels: 8, out_channels: 16, freq_dim: 4, time_steps: 5\n",
      "in_channels: 16, out_channels: 32, freq_dim: 1, time_steps: 1\n",
      "\n",
      "Feature keys: ['signal']\n",
      "Label key: ['label']\n",
      "Model mode: multiclass\n",
      "Total parameters: 25,326\n",
      "Trainable parameters: 25,326\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.models import ContraWR\n",
    "\n",
    "model = ContraWR(\n",
    "    dataset=sample_dataset,\n",
    ").to(device)\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f\"Feature keys: {model.feature_keys}\")\n",
    "print(f\"Label key: {model.label_keys}\")\n",
    "print(f\"Model mode: {model.mode}\")\n",
    "print(f\"Total parameters: {total_params:,}\")\n",
    "print(f\"Trainable parameters: {trainable_params:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35201769",
   "metadata": {},
   "source": [
    "# 7. Test Forward Pass\n",
    "Verify the model can process a batch and compute the contrastive loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4bfd1d77",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model output keys: dict_keys(['loss', 'y_prob', 'y_true', 'logit'])\n",
      "Loss value: 2.2494\n",
      "Output probability shape: torch.Size([32, 6])\n"
     ]
    }
   ],
   "source": [
    "# Move batch to device\n",
    "test_batch = {key: value.to(device) if hasattr(value, 'to') else value \n",
    "              for key, value in first_batch.items()}\n",
    "\n",
    "# Forward pass\n",
    "with torch.no_grad():\n",
    "    output = model(**test_batch)\n",
    "\n",
    "print(\"Model output keys:\", output.keys())\n",
    "if 'loss' in output:\n",
    "    print(f\"Loss value: {output['loss'].item():.4f}\")\n",
    "if 'y_prob' in output:\n",
    "    print(f\"Output probability shape: {output['y_prob'].shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "693e383f",
   "metadata": {},
   "source": [
    "# 8. Configure Trainer\n",
    "Wrap the model with the PyHealth Trainer and define optimization hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cf9a1369",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ContraWR(\n",
      "  (encoder): Sequential(\n",
      "    (0): ResBlock2D(\n",
      "      (conv1): Conv2d(7, 8, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (bn1): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ELU(alpha=1.0)\n",
      "      (conv2): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (bn2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "      (downsampler): Sequential(\n",
      "        (0): Conv2d(7, 8, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (1): ResBlock2D(\n",
      "      (conv1): Conv2d(8, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ELU(alpha=1.0)\n",
      "      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "      (downsampler): Sequential(\n",
      "        (0): Conv2d(8, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "    (2): ResBlock2D(\n",
      "      (conv1): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ELU(alpha=1.0)\n",
      "      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "      (downsampler): Sequential(\n",
      "        (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "    )\n",
      "  )\n",
      "  (fc): Linear(in_features=32, out_features=6, bias=True)\n",
      ")\n",
      "Metrics: None\n",
      "Device: cpu\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.trainer import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    device=str(device),\n",
    "    enable_logging=False,\n",
    ")\n",
    "\n",
    "training_config = {\n",
    "    \"epochs\": 10,\n",
    "    \"optimizer_params\": {\"lr\": 1e-3},\n",
    "    \"max_grad_norm\": 5.0,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea8b2908",
   "metadata": {},
   "source": [
    "# 9. Train the Model\n",
    "Launch the self-supervised training loop to learn representations from sleep signal data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "644a437f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training:\n",
      "Batch size: 32\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: 5.0\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x0000029B647EE2F0>\n",
      "Monitor: None\n",
      "Monitor criterion: max\n",
      "Epochs: 10\n",
      "Patience: None\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 / 10: 100%|██████████| 120/120 [00:06<00:00, 18.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-0, step-120 ---\n",
      "loss: 0.9076\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 36.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-0, step-120 ---\n",
      "accuracy: 0.7646\n",
      "f1_macro: 0.2430\n",
      "f1_micro: 0.7646\n",
      "loss: 0.6233\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1 / 10: 100%|██████████| 120/120 [00:06<00:00, 18.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-1, step-240 ---\n",
      "loss: 0.6502\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 36.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-1, step-240 ---\n",
      "accuracy: 0.7737\n",
      "f1_macro: 0.2510\n",
      "f1_micro: 0.7737\n",
      "loss: 0.6300\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2 / 10: 100%|██████████| 120/120 [00:06<00:00, 19.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-2, step-360 ---\n",
      "loss: 0.6182\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 33.93it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-2, step-360 ---\n",
      "accuracy: 0.7810\n",
      "f1_macro: 0.2533\n",
      "f1_micro: 0.7810\n",
      "loss: 0.6012\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3 / 10: 100%|██████████| 120/120 [00:06<00:00, 19.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-3, step-480 ---\n",
      "loss: 0.5995\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 37.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-3, step-480 ---\n",
      "accuracy: 0.7755\n",
      "f1_macro: 0.2516\n",
      "f1_micro: 0.7755\n",
      "loss: 0.5852\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4 / 10: 100%|██████████| 120/120 [00:06<00:00, 19.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-4, step-600 ---\n",
      "loss: 0.5886\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 36.99it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-4, step-600 ---\n",
      "accuracy: 0.7974\n",
      "f1_macro: 0.2598\n",
      "f1_micro: 0.7974\n",
      "loss: 0.5643\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 5 / 10: 100%|██████████| 120/120 [00:05<00:00, 20.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-5, step-720 ---\n",
      "loss: 0.5647\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 38.63it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-5, step-720 ---\n",
      "accuracy: 0.7920\n",
      "f1_macro: 0.2941\n",
      "f1_micro: 0.7920\n",
      "loss: 0.5486\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 6 / 10: 100%|██████████| 120/120 [00:06<00:00, 19.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-6, step-840 ---\n",
      "loss: 0.5616\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 25.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-6, step-840 ---\n",
      "accuracy: 0.8011\n",
      "f1_macro: 0.2607\n",
      "f1_micro: 0.8011\n",
      "loss: 0.5348\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 7 / 10: 100%|██████████| 120/120 [00:07<00:00, 15.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-7, step-960 ---\n",
      "loss: 0.5567\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 33.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-7, step-960 ---\n",
      "accuracy: 0.7956\n",
      "f1_macro: 0.2559\n",
      "f1_micro: 0.7956\n",
      "loss: 0.5355\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 8 / 10: 100%|██████████| 120/120 [00:05<00:00, 20.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-8, step-1080 ---\n",
      "loss: 0.5500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 35.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-8, step-1080 ---\n",
      "accuracy: 0.7920\n",
      "f1_macro: 0.3068\n",
      "f1_micro: 0.7920\n",
      "loss: 0.5316\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 9 / 10: 100%|██████████| 120/120 [00:06<00:00, 19.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-9, step-1200 ---\n",
      "loss: 0.5496\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 18/18 [00:00<00:00, 38.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-9, step-1200 ---\n",
      "accuracy: 0.7920\n",
      "f1_macro: 0.3132\n",
      "f1_micro: 0.7920\n",
      "loss: 0.5272\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "trainer.train(\n",
    "    train_dataloader=train_loader,\n",
    "    val_dataloader=val_loader,\n",
    "    **training_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ea42dbf",
   "metadata": {},
   "source": [
    "# 10. Save Model (Optional)\n",
    "Save the trained model for future use or fine-tuning on downstream tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d34bd11",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model saved to: contrawr_sleepedf_model.pth\n"
     ]
    }
   ],
   "source": [
    "# Save model checkpoint\n",
    "save_path = \"contrawr_sleepedf_model.pth\"\n",
    "torch.save({\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'config': training_config,\n",
    "}, save_path)\n",
    "print(f\"Model saved to: {save_path}\")\n",
    "\n",
    "# To load the model later:\n",
    "# checkpoint = torch.load(save_path)\n",
    "# model.load_state_dict(checkpoint['model_state_dict'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.10.11)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
